import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QTableWidget, QTableWidgetItem, QVBoxLayout, QWidget, QPushButton, QFileDialog, QInputDialog
from PyQt5.QtGui import QColor
from datetime import datetime, timedelta
import pandas as pd

class ScheduleApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.shift_options = ["Morning", "Afternoon", "Evening", "Off"]  # 班次选项
        self.initUI()

    def initUI(self):
        self.setWindowTitle('排班表')
        self.setGeometry(100, 100, 800, 600)

        # 排班周期为7天
        self.schedule_days = 7

        # 计算本周的星期一的日期
        today = datetime.now().date()
        self.start_date = today - timedelta(days=today.weekday())  # weekday()返回0表示星期一

        # 定义一个字典用于转换星期几为中文
        self.weekday_map = {
            'Monday': '星期一', 'Tuesday': '星期二',
            'Wednesday': '星期三', 'Thursday': '星期四',
            'Friday': '星期五', 'Saturday': '星期六',
            'Sunday': '星期日'
        }

        # 生成日期标签
        date_labels = [
            f"{self.start_date + timedelta(days=i):%Y/%#m/%#d} " +  # 使用指定格式并去除前导零
            self.weekday_map[(self.start_date + timedelta(days=i)).strftime('%A')]
            for i in range(self.schedule_days)
        ]

        # 如果运行在不支持 %#m 或 %#d 的平台上，则使用以下替代方案：
        if sys.platform == "win32":
            date_labels = [
                f"{self.start_date + timedelta(days=i):%Y/%m/%d}".replace("/0", "/").replace("/", "/ ") +
                self.weekday_map[(self.start_date + timedelta(days=i)).strftime('%A')]
                for i in range(self.schedule_days)
            ]

        # 创建表格（姓名列 +7天排班）
        self.table = QTableWidget()
        self.table.setColumnCount(self.schedule_days + 1)
        self.table.setHorizontalHeaderLabels(['姓名'] + date_labels)

        # 示例数据（每行包含姓名和7天的排班项）
        data = [
            ("Alice", "Morning", "Afternoon", "Off", "Morning", "Afternoon", "Evening", "Off"),
            ("Bob", "Afternoon", "Off", "Morning", "Afternoon", "Evening", "Morning", "Off"),
            ("Charlie", "Off", "Morning", "Afternoon", "Off", "Morning", "Afternoon", "Evening"),
        ]

        self.table.setRowCount(len(data))
        for row_index, row_data in enumerate(data):
            # 设置姓名列（列0）
            name = row_data[0]
            self.table.setItem(row_index, 0, QTableWidgetItem(name))

            # 设置排班列（列1 到列7）
            for day in range(1, self.schedule_days + 1):  # 遍历 1-7 天的排班数据
                schedule = row_data[day] if day < len(row_data) else "N/A"
                item = QTableWidgetItem(schedule)
                self.updateItemColor(item)
                self.table.setItem(row_index, day, item)

        # 设置双击事件监听器
        self.table.itemDoubleClicked.connect(self.selectShift)

        # 导出按钮
        self.export_button = QPushButton('导出到Excel', self)
        self.export_button.clicked.connect(self.export_to_excel)

        # 布局
        layout = QVBoxLayout()
        layout.addWidget(self.table)
        layout.addWidget(self.export_button)

        container = QWidget()
        container.setLayout(layout)
        self.setCentralWidget(container)

    def selectShift(self, item):
        """
        当用户双击非标题单元格时调用此函数，允许用户选择班次并更新单元格内容。
        """
        print(f"Item clicked at row: {item.row()}, column: {item.column()}")  # 调试信息
        if item.row() >= 0:  # 只处理数据行，排除标题行
            current_shift = item.text()
            shift, ok = QInputDialog.getItem(self, "选择班次", "选择班次:", self.shift_options,
                                             self.shift_options.index(current_shift) if current_shift in self.shift_options else 0, False)
            if ok and shift:
                item.setText(shift)
                self.updateItemColor(item)
        else:
            print("Title row clicked, ignoring.")  # 调试信息

    def updateItemColor(self, item):
        """
        根据单元格的班次内容更新单元格的背景颜色。
        """
        if item.text() == "Off":
            item.setBackground(QColor(255, 0, 0))  # 设置为红色
        else:
            item.setBackground(QColor(255, 255, 255))  # 默认白色

    def export_to_excel(self):
        rows = self.table.rowCount()
        cols = self.table.columnCount()
        data = []

        # 提取表格数据
        for row in range(rows):
            row_data = [
                self.table.item(row, col).text() if self.table.item(row, col) else ""
                for col in range(cols)
            ]
            data.append(row_data)

        # 创建DataFrame时，为确保第一行（标题行）使用正确的日期格式，
        # 我们手动构建标题行，使用之前定义的date_labels
        headers = ['姓名'] + [
            f"{self.start_date + timedelta(days=i):%Y/%#m/%#d}".replace("/0", "/").replace("/", "/ ") +
            self.weekday_map[(self.start_date + timedelta(days=i)).strftime('%A')]
            for i in range(self.schedule_days)
        ]

        df = pd.DataFrame(data, columns=headers)

        file_path, _ = QFileDialog.getSaveFileName(self, "保存到Excel", "", "Excel文件 (*.xlsx)")
        if file_path:
            df.to_excel(file_path, index=False)
            print(f"已保存排班表到: {file_path}")

if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = ScheduleApp()
    ex.show()
    sys.exit(app.exec())